package top.haijunit.sample.navi.utils.stream;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Set;
import java.util.function.*;
import java.util.stream.Collector;

import static java.util.Objects.requireNonNull;

/**
 * 批量收集和处理数据
 *
 * @author zhanghaijun
 * @date 2024/7/8 18:36
 * @description 批量收集和处理数据
 */
public class BatchCollector<T> implements Collector<T, List<T>, List<T>> {

    // 批次大小，用于指定每次处理多少元素。
    private final int batchSize;
    // 批处理器，用于处理每个批次的元素。
    private final Consumer<List<T>> batchProcessor;

    public BatchCollector(int batchSize, Consumer<List<T>> batchProcessor) {
        requireNonNull(batchProcessor, "Batch processor cannot be null");
        if (batchSize <= 0) {
            throw new IllegalArgumentException("Batch size must be greater than 0");
        }
        this.batchSize = batchSize;
        this.batchProcessor = batchProcessor;
    }

    /**
     * 累积容器
     */
    @Override
    public Supplier<List<T>> supplier() {
        return ArrayList::new;
    }

    /**
     * 如何将单个元素添加到累积容器中，并在达到批处理大小时处理批量数据。
     */
    @Override
    public BiConsumer<List<T>, T> accumulator() {
        return (ts, t) -> {
            ts.add(t);
            if (ts.size() >= batchSize) {
                batchProcessor.accept(new ArrayList<>(ts));
                ts.clear();
            }
        };
    }

    /**
     * 如何合并两个部分累积的结果。这在并行流中非常重要。当合并的结果达到批处理大小时，立即处理批量数据。
     */
    @Override
    public BinaryOperator<List<T>> combiner() {
        return (ts, ots) -> {
            ts.addAll(ots);
            if (ts.size() >= batchSize) {
                batchProcessor.accept(new ArrayList<>(ts));
                ts.clear();
            }
            return ts;
        };
    }

    /**
     * 累积过程结束时被调用。如果有剩余未处理的数据，在此方法中处理这些数据。
     */
    @Override
    public Function<List<T>, List<T>> finisher() {
        return ts -> {
            if (!ts.isEmpty()) {
                batchProcessor.accept(new ArrayList<>(ts));
            }
            return Collections.emptyList();
        };
    }

    /**
     * 返回一个空的特性集合。特性用于描述这个收集器的行为。
     */
    @Override
    public Set<Characteristics> characteristics() {
        return Collections.emptySet();
    }
}
